
from stable_diffusion.ddpm import DDPM
import matplotlib.pyplot as plt
from gdatasets import beta_schedule, alpha_bar, alpha
import torch
import numpy as np

def beta_schedule(step_num):
    start = 1e-4
    end = 0.02

    return np.linspace(start, end, step_num)

device = torch.device("cuda")

r = 10
c = 10
batch_size = r * c
size = (batch_size, 3, 32, 32)


model = DDPM().to(device)
model.load_state_dict(torch.load("./best_model", weights_only=True, map_location=device))
model.eval()

beta = beta_schedule(1000)
alpha_val = alpha(beta)
alpha_bar_val = alpha_bar(beta)

beta = torch.tensor(beta, dtype=torch.float32, device=device)
alpha_val = torch.tensor(alpha_val, dtype=torch.float32, device=device)
alpha_bar_val = torch.tensor(alpha_bar_val, dtype=torch.float32, device=device)


with torch.inference_mode():
    x = torch.normal(0, 1, size = size, dtype=torch.float32, device=device)
    for i in range(1000):
        t = 1000 - i
        if t == 1:
            z = torch.zeros(size=size, dtype=torch.float32, device=device)
        else:
            z = torch.normal(0, 1, size=size, dtype=torch.float32, device=device)

        param1 = 1 / torch.sqrt(alpha_val[t - 1])
        param2 = (1 - alpha_val[t - 1]) / (torch.sqrt(1 - alpha_bar_val[t - 1]))
        gamma = torch.sqrt(beta[t - 1])

        t_tensor = torch.full((batch_size, 1), t, dtype=torch.float32, device=device)
        x = param1 *(x - param2 * model(x, t_tensor).reshape(-1, 3, 32, 32)) + gamma * z


x = x.detach().cpu().numpy().transpose(0, 2, 3, 1)

fig, ax = plt.subplots(r, c)
for i in range(r):
    for j in range(c):
        idx = i * c + j
        ax[i, j].imshow(x[idx])
        ax[i, j].axis("off")


plt.show()
